DQN with CartPole

import gymnasium as gym
import math
import random
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple, deque
from itertools import count

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

env = gym.make("CartPole-v1")

# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
    from IPython import display

plt.ion()

# if GPU is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Set up

• \(r = 1\) for every incremental timestep.
• \(s_T = \text{2.4 units away from center or falls over too far}\)

env = gym.make("CartPole-v1")

device = "cuda" if torch.cuda.is_available() else "cpu"
device

'cuda'

Transition = namedtuple("Transition",("state","action","next_state","reward"))

class ReplayMemory():
    def __init__(self,capacity):
        self.memory = deque([],maxlen = capacity)
    
    def push(self,*args):
        """transition save"""
        self.memory.append(Transition(*args))
    
    def sample(self,batch_size):
        """
        Input : batch_size
        Output : randomly sampled transitions in memory
        """
        return random.sample(self.memory,batch_size)
    
    def __len__(self):
        return len(self.memory)
        

Q-network

class DQN(nn.Module):
    def __init__(self,n_obs,n_actions):
        super().__init__()
        self.layer1 = nn.Linear(n_obs,128)
        self.layer2 = nn.Linear(128,64)
        self.layer3 = nn.Linear(64,n_actions)
        self.activation = nn.ReLU()
    def forward(self,x):
        l1_out = self.activation(self.layer1(x))
        l2_out = self.activation(self.layer2(l1_out))
        return self.layer3(l2_out)

dqn = DQN(32,10)
_t = torch.randn(10,32)
dqn(_t)[0]

tensor([-0.0175, -0.0064, -0.0310,  0.0675, -0.0197,  0.2433,  0.1908, -0.0816,
        -0.0720, -0.1309], grad_fn=<SelectBackward0>)

# BATCH_SIZE is the number of transitions sampled from the replay buffer
# GAMMA is the discount factor as mentioned in the previous section
# EPS_START is the starting value of epsilon
# EPS_END is the final value of epsilon
# EPS_DECAY controls the rate of exponential decay of epsilon, higher means a slower decay
# TAU is the update rate of the target network
# LR is the learning rate of the ``AdamW`` optimizer
BATCH_SIZE = 128
GAMMA = 0.99
EPS_START = 0.9
EPS_END = 0.05
EPS_DECAY = 1000
TAU = 0.005
LR = 1e-4

# Get number of actions from gym action space
n_actions = env.action_space.n
# Get the number of state observations
state, info = env.reset()
n_observations = len(state)

policy_net = DQN(n_observations, n_actions).to(device)
target_net = DQN(n_observations, n_actions).to(device)
target_net.load_state_dict(policy_net.state_dict())

optimizer = optim.AdamW(policy_net.parameters(), lr=LR, amsgrad=True)
memory = ReplayMemory(10000)


steps_done = 0


def select_action(state):
    global steps_done
    sample = random.random()
    eps_threshold = EPS_END + (EPS_START - EPS_END) * \
        math.exp(-1. * steps_done / EPS_DECAY)
    steps_done += 1
    if sample > eps_threshold:
        with torch.no_grad():
            # t.max(1) will return the largest column value of each row.
            # second column on max result is index of where max element was
            # found, so we pick action with the larger expected reward.
            return policy_net(state).max(1)[1].view(1, 1)
    else:
        return torch.tensor([[env.action_space.sample()]], device=device, dtype=torch.long)


episode_durations = []


def plot_durations(show_result=False):
    plt.figure(1)
    durations_t = torch.tensor(episode_durations, dtype=torch.float)
    if show_result:
        plt.title('Result')
    else:
        plt.clf()
        plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Duration')
    plt.plot(durations_t.numpy())
    # Take 100 episode averages and plot them too
    if len(durations_t) >= 100:
        means = durations_t.unfold(0, 100, 1).mean(1).view(-1)
        means = torch.cat((torch.zeros(99), means))
        plt.plot(means.numpy())

    plt.pause(0.001)  # pause a bit so that plots are updated
    if is_ipython:
        if not show_result:
            display.display(plt.gcf())
            display.clear_output(wait=True)
        else:
            display.display(plt.gcf())

steps_done = 0

def select_action(state):
    global steps_done
    sample = random.random()
    eps_threshold = EPS_END + (EPS_START - EPS_END) * math.exp(-1 * steps_done / EPS_DECAY) #ϵ의 상한,하한을 만드는 전략,이거써야겠다.
    steps_done += 1
    """
    if steps_done%10 == 0:
        print(eps_threshold,steps_done)
    """
    #ϵ-greedy policy
    if sample > eps_threshold:
        with torch.no_grad():
            return policy_net(state).max(1)[1].view(1,1)
    else:
        return torch.tensor([[env.action_space.sample()]],device=device,dtype = torch.long)
    

_state = torch.randn(2,4).to(device)
print(_state)
_q_out = policy_net(_state)
print(_q_out)
_max_q = policy_net(_state).max(axis=1)
print(_max_q)

tensor([[ 0.4280,  0.3374, -0.6418,  0.9572],
        [ 0.1284,  1.3429,  0.0458,  0.7359]], device='cuda:0')
tensor([[ 0.0939, -0.0166],
        [ 0.0718,  0.0492]], device='cuda:0', grad_fn=<AddmmBackward0>)
torch.return_types.max(
values=tensor([0.0939, 0.0718], device='cuda:0', grad_fn=<MaxBackward0>),
indices=tensor([0, 0], device='cuda:0'))

episode_durations = []
def plot_durations(show_result=False):
    plt.figure(1)
    durations_t = torch.tensor(episode_durations, dtype=torch.float)
    if show_result:
        plt.title('Result')
    else:
        plt.clf()
        plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Duration')
    plt.plot(durations_t.numpy())
    # Take 100 episode averages and plot them too
    if len(durations_t) >= 100:
        means = durations_t.unfold(0, 100, 1).mean(1).view(-1)
        means = torch.cat((torch.zeros(99), means))
        plt.plot(means.numpy())

    plt.pause(0.001)  # pause a bit so that plots are updated
    if is_ipython:
        if not show_result:
            display.display(plt.gcf())
            display.clear_output(wait=True)
        else:
            display.display(plt.gcf())

• optimize_model functio n
  1. samples a batch \(\sim \mathcal{D}\)
  2. concatenate all the tensors into a single one
  3. compute Q and combines them into loss

def optimize_model():
    if len(memory) < BATCH_SIZE:
        # BATCH_SIZE보다 작을 경우 아예 다음 코드가 실행이 안됨.
        # 즉,Experience replay가 안됨.
        return
    global transitions,batch,_t,non_final_mask,non_final_next_states,state_batch\
    ,reward_batch,action_batch,state_action_values,non_final_next_states,__t,before_masked,after_masked\
    ,expected_state_action_values
    transitions = memory.sample(BATCH_SIZE)
    # 메모리에서 BATCH_SIZE만큼의 Transition들을 가져옴    
    _t = zip(*transitions)
    batch = Transition(*zip(*transitions))
    # Transition이라는 named tuple에 BATCH_SIZE만큼의 state,action ...를 따로따로 저장.
    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                          batch.next_state)), device=device, dtype=torch.bool)
    non_final_next_states = torch.cat([s for s in batch.next_state
                                                if s is not None])
    # next_state가 Terminal이 아닌 state값들만 가져옴
    # Terminal이 None이므로 여기서는 None값을 제외한 next_state들을 가져옴
    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)
    
    # poicy_net은 현재 s_t,a_t에 대한 Q_value임.
    # 일단 현재 t에 대한 Q_value를 계산!
    #__t = policy_net(state_batch)
    state_action_values = policy_net(state_batch).gather(1, action_batch)
    """
    next_state_value를 계산하는 방법
    #1. terminal_state를 0 non_terminal_state를 1로 놓은 boolean tensor를 만든다.(mask)
    #2. 모든 state에 대한 value function인 target network를 가져오고
    #3. target network에서 mask를 활용하여 final state인 것은 value를 0으로 계산한다.
    # 이렇게 함으로 임의의 파라미터로 인한 target network가 Terminal state에 대한 value를 0이아닌 값을 주는 것을 막을 수 있다.
    """
    next_state_values = torch.zeros(BATCH_SIZE, device=device)
    #before_masked = torch.zeros(BATCH_SIZE, device=device)
    with torch.no_grad():
        next_state_values[non_final_mask] = target_net(non_final_next_states).max(1)[0]
    # Compute the expected Q values
    # after_masked = next_state_values
    expected_state_action_values = (next_state_values * GAMMA) + reward_batch

    # Compute Huber loss
    criterion = nn.SmoothL1Loss()
    loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))

    # Optimize the model
    optimizer.zero_grad()
    loss.backward()
    # In-place gradient clipping
    torch.nn.utils.clip_grad_value_(policy_net.parameters(), 100)
    optimizer.step()

if torch.cuda.is_available():
    num_episodes = 500
else:
    num_episodes = 50

for i_episode in range(num_episodes):
    # Initialize the environment and get it's state
    state, info = env.reset()
    state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
    for t in count():
        action = select_action(state)
        observation, reward, terminated, truncated, _ = env.step(action.item())
        reward = torch.tensor([reward], device=device)
        done = terminated or truncated

        if terminated:
            next_state = None
        else:
            next_state = torch.tensor(observation, dtype=torch.float32, device=device).unsqueeze(0)

        # Store the transition in memory
        memory.push(state, action, next_state, reward)

        # Move to the next state
        state = next_state

        # Perform one step of the optimization (on the policy network)
        optimize_model()

        # Soft update of the target network's weights
        # θ′ ← τ θ + (1 −τ )θ′
        target_net_state_dict = target_net.state_dict()
        policy_net_state_dict = policy_net.state_dict()
        for key in policy_net_state_dict:
            target_net_state_dict[key] = policy_net_state_dict[key]*TAU + target_net_state_dict[key]*(1-TAU)
        target_net.load_state_dict(target_net_state_dict)

        if done:
            episode_durations.append(t + 1)
            plot_durations()
            break

print('Complete')
plot_durations(show_result=True)
plt.ioff()
plt.show()

Complete

<Figure size 640x480 with 0 Axes>

<Figure size 640x480 with 0 Axes>

s,_ = env.reset()
policy_net(torch.tensor(s).to(device))

tensor([99.5481, 99.2072], device='cuda:0', grad_fn=<AddBackward0>)

코드 정리

unpacking

def intro(name,age,univeristy_department):
    print(name,age,univeristy_department)

info = ["신호연",14,["HY","AI-semiconductor"]]

intro(info)

TypeError: intro() missing 2 required positional arguments: 'age' and 'univeristy_department'

• 함수는 3개의 파라미터에 대한 값(인자)를 필요로하는데 넣어주는 건 하나여서 오류 발생
• 리스트의 3개의 요소가 각각의 파라미터에 들어가는 인자가 되어야 함.
• 이때 unpacking, 즉 iterable한 객체안에 있는 요소를 input으로 넣어주는 역할

intro(*info)

신호연 14 ['HY', 'AI-semiconductor']

packing(*args , **kwargs)

def hello(*args):
    print(args)
hello(1,2.5,"hi",True)

(1, 2.5, 'hi', True)

def hello(**kwargs):
    print(kwargs)

hello(kwargs1 = "g",kw = 2,age = 26)

{'kwargs1': 'g', 'kw': 2, 'age': 26}

torch.gather

t = torch.tensor([i+1 for i in range(9)]).reshape(3,3)
index = torch.tensor([[0,1,2],[1,2,0]])
torch.gather(t,dim=0,index=index)

tensor([[1, 5, 9],
        [4, 8, 3]])

t

tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])

t = torch.tensor([i+1 for i in range(9)]).reshape(3,3)
index = torch.tensor([[0,1],[1,2]])
torch.gather(t,dim=1,index=index)

tensor([[1, 2],
        [5, 6]])

[[1,2,]
,[5,6]
,[9,7]]

[[1, 2], [5, 6], [9, 7]]

torch.gather(t,dim=1,index=torch.tensor([[1,1]]))

tensor([[2, 2]])

t = torch.tensor([i for i in range(100)]).reshape(10,10)
index = torch.tensor([i for i in range(10)]*2).reshape(10,2)[:-3]
g_t = torch.gather(t,1,index)
print(t.shape,index.shape)
print(g_t)

torch.Size([10, 10]) torch.Size([7, 2])
tensor([[ 0,  1],
        [12, 13],
        [24, 25],
        [36, 37],
        [48, 49],
        [50, 51],
        [62, 63]])

A = torch.tensor([i+1 for i in range(4)]).reshape(2,2)
torch.gather(A,index=index,dim=1).reshape(1,2)
index = torch.tensor([[0],[1]])
A.dim(),index.dim(),index,A.shape,index.shape

(2,
 2,
 tensor([[0],
         [1]]),
 torch.Size([2, 2]),
 torch.Size([2, 1]))

map()

# 리스트에 값을 하나씩 더해서 새로운 리스트를 만드는 작업
myList = [1, 2, 3, 4, 5]

# for 반복문 이용
result1 = []
for val in myList:
    result1.append(val + 1)

print(f'result1 : {result1}')


# map 함수 이용
def add_one(n):
    return n + 1


result2 = list(map(add_one, myList))  # map반환을 list 로 변환
print(f'result2 : {result2}')

result1 : [2, 3, 4, 5, 6]
result2 : [2, 3, 4, 5, 6]

myList

result1 = []
for val in myList:
    result1.append(val+1)
print(result1)

[2, 3, 4, 5, 6]

result2 = list(map(lambda x : x+1,myList))
result2

[2, 3, 4, 5, 6]

zip()

a = ["1",2]
b = [3,"5"]
c = [2,4]
for t in zip(a,b,c):
    print(t)

('1', 3, 2)
(2, '5', 4)

torch.unsuqeeuze()

import torch
ind_A  = torch.tensor([1,0,3])
ind_A.unsqueeze(1).unsqueeze(1) == ind_A.unsqueeze(1).unsqueeze(2)

tensor([[[True]],

        [[True]],

        [[True]]])

torch.tensor([i for i in range(8)]).reshape(2,4)

tensor([[0, 1, 2, 3],
        [4, 5, 6, 7]])

check

len(transitions)

128

before_masked

tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0.], device='cuda:0')

after_masked

tensor([0.4509, 0.6724, 0.6550, 0.3936, 0.4086, 0.3233, 0.3974, 0.6300, 0.0000,
        0.3620, 1.1384, 0.3229, 0.4893, 0.9286, 0.3843, 0.4094, 0.3200, 0.6750,
        0.3486, 0.4730, 0.7877, 0.5478, 0.3448, 0.0000, 0.5176, 0.3613, 0.3535,
        0.3946, 0.8328, 0.3183, 0.4061, 0.3919, 0.3045, 0.3789, 0.7998, 0.8365,
        0.9102, 0.3106, 0.3463, 0.4299, 0.3520, 0.6124, 0.3324, 0.6185, 0.4004,
        0.3432, 0.6711, 0.3977, 0.3285, 0.6386, 0.3412, 0.7142, 0.5016, 0.3459,
        0.0000, 0.3325, 0.3951, 0.3529, 0.4550, 0.5403, 0.3315, 0.3641, 0.3638,
        0.6286, 0.3520, 0.4379, 0.3301, 0.7780, 0.3207, 0.4688, 0.0000, 0.4506,
        0.3452, 0.4396, 0.6683, 0.3663, 0.8068, 0.8672, 0.6586, 0.3427, 0.4015,
        0.3412, 0.6464, 0.8239, 0.3989, 0.3136, 1.3251, 0.5879, 0.3562, 0.3875,
        0.4905, 0.3174, 0.3465, 0.3737, 1.2304, 0.4512, 0.3386, 0.7941, 0.3301,
        0.3950, 0.3471, 0.3744, 1.0437, 0.3394, 0.3246, 0.5522, 0.3403, 0.4993,
        0.4461, 0.5644, 0.6235, 0.4678, 0.7648, 0.3878, 0.3846, 0.4321, 1.1868,
        0.3423, 0.4125, 0.3476, 0.6812, 0.0000, 0.6112, 0.5141, 0.3533, 0.3334,
        0.3472, 0.2959], device='cuda:0')

t = Transition(*zip(*transitions))

t.next_state

(tensor([[ 0.1073,  1.7679, -0.1467, -2.5517]], device='cuda:0'),
 tensor([[ 0.0853,  1.2131, -0.1016, -1.8380]], device='cuda:0'),
 tensor([[-0.0783,  0.5542,  0.1493, -0.4089]], device='cuda:0'),
 tensor([[ 0.3524,  1.9037, -0.1091, -2.0654]], device='cuda:0'),
 tensor([[ 0.0199,  0.1916, -0.0173, -0.3075]], device='cuda:0'),
 tensor([[-0.0198,  0.4050,  0.0023, -0.5490]], device='cuda:0'),
 tensor([[ 0.0097,  0.3763,  0.0113, -0.3715]], device='cuda:0'),
 tensor([[-0.0265,  0.5466,  0.1104, -0.2387]], device='cuda:0'),
 tensor([[ 0.0936,  1.0071, -0.1436, -1.5632]], device='cuda:0'),
 tensor([[ 0.0229,  0.4268, -0.0829, -0.7092]], device='cuda:0'),
 tensor([[-0.0352, -0.0211, -0.0086, -0.0348]], device='cuda:0'),
 tensor([[ 0.0347,  0.6346, -0.1449, -1.1158]], device='cuda:0'),
 tensor([[ 0.1541,  0.8217, -0.1321, -1.1422]], device='cuda:0'),
 tensor([[ 0.1673,  1.3828, -0.1596, -2.1547]], device='cuda:0'),
 tensor([[-0.0433,  0.4303, -0.0228, -0.6093]], device='cuda:0'),
 tensor([[-0.0285,  0.5732,  0.0584, -0.7044]], device='cuda:0'),
 tensor([[-0.0856,  0.3615,  0.1526, -0.1680]], device='cuda:0'),
 tensor([[ 0.2314,  1.5120,  0.0065, -1.4392]], device='cuda:0'),
 tensor([[ 0.0240, -0.1551, -0.0162,  0.2688]], device='cuda:0'),
 tensor([[ 0.0085,  0.0406,  0.0099, -0.0364]], device='cuda:0'),
 tensor([[-0.0609,  0.1449,  0.1419, -0.0150]], device='cuda:0'),
 tensor([[-0.0229,  0.7601, -0.0386, -1.2230]], device='cuda:0'),
 tensor([[ 0.0440,  0.9616, -0.0284, -1.2477]], device='cuda:0'),
 tensor([[-0.0254, -0.0371,  0.1096,  0.2920]], device='cuda:0'),
 tensor([[ 0.0575,  0.5710, -0.0283, -0.8761]], device='cuda:0'),
 None,
 tensor([[-0.0227, -0.2063,  0.0239,  0.3246]], device='cuda:0'),
 tensor([[-0.0155, -0.5877,  0.0317,  0.8373]], device='cuda:0'),
 tensor([[ 0.2050,  1.3173,  0.0296, -1.1560]], device='cuda:0'),
 tensor([[-1.4073e-02, -3.7479e-01,  2.4043e-04,  6.0724e-01]], device='cuda:0'),
 tensor([[ 0.0105,  1.1017,  0.1021, -1.0711]], device='cuda:0'),
 tensor([[-0.0261, -0.2336,  0.1154,  0.6171]], device='cuda:0'),
 tensor([[ 6.1860e-04, -6.2460e-01,  4.1290e-02,  9.1309e-01]], device='cuda:0'),
 tensor([[ 0.0380, -0.0239,  0.0852,  0.3501]], device='cuda:0'),
 tensor([[ 0.0263, -0.1556, -0.0214,  0.2805]], device='cuda:0'),
 tensor([[-0.0045,  0.0275, -0.0069,  0.0019]], device='cuda:0'),
 tensor([[ 0.0352, -0.2067,  0.0393,  0.3709]], device='cuda:0'),
 tensor([[ 0.0520,  0.8176,  0.0014, -1.0424]], device='cuda:0'),
 tensor([[ 0.0507,  0.7092,  0.0657, -0.4313]], device='cuda:0'),
 tensor([[ 0.0496,  0.2426, -0.0386, -0.2745]], device='cuda:0'),
 tensor([[ 0.0365, -0.0180,  0.0587,  0.2197]], device='cuda:0'),
 tensor([[-0.0416,  0.3695, -0.0076, -0.6290]], device='cuda:0'),
 tensor([[-0.0201, -0.0344,  0.0939,  0.2313]], device='cuda:0'),
 tensor([[ 0.0504,  0.5380,  0.0743, -0.0478]], device='cuda:0'),
 tensor([[-0.0088,  0.4327, -0.0728, -0.6648]], device='cuda:0'),
 tensor([[-0.0349,  0.2070,  0.0439, -0.1921]], device='cuda:0'),
 tensor([[-0.0885, -0.0239,  0.1448,  0.3172]], device='cuda:0'),
 tensor([[-0.0222,  0.4311, -0.0532, -0.6276]], device='cuda:0'),
 tensor([[ 0.3219,  1.5108, -0.1147, -1.4554]], device='cuda:0'),
 tensor([[ 0.0611,  0.7320,  0.0734, -0.3161]], device='cuda:0'),
 tensor([[ 0.1906,  1.3128, -0.0039, -1.0927]], device='cuda:0'),
 tensor([[-0.0327, -0.1801,  0.0309,  0.3245]], device='cuda:0'),
 tensor([[-0.0105, -0.1798, -0.0061,  0.3165]], device='cuda:0'),
 tensor([[ 0.0474,  0.8313, -0.1672, -1.4501]], device='cuda:0'),
 tensor([[ 0.0024,  0.7370,  0.1039, -0.4280]], device='cuda:0'),
 tensor([[-2.4166e-04, -2.2275e-01,  5.1142e-02,  3.7374e-01]], device='cuda:0'),
 tensor([[-0.0464, -0.4346,  0.1129,  0.7364]], device='cuda:0'),
 tensor([[ 0.0681,  0.4396, -0.0616, -0.6103]], device='cuda:0'),
 tensor([[-0.0073,  0.0277, -0.0010, -0.0024]], device='cuda:0'),
 tensor([[ 0.2655,  1.5088, -0.0638, -1.4074]], device='cuda:0'),
 tensor([[-0.0523,  0.5501,  0.1281, -0.3174]], device='cuda:0'),
 tensor([[-0.0268, -0.4017,  0.0304,  0.6248]], device='cuda:0'),
 tensor([[-0.0252, -0.3750,  0.0187,  0.6112]], device='cuda:0'),
 tensor([[ 0.0380,  1.1537, -0.1385, -1.8957]], device='cuda:0'),
 tensor([[ 0.0393, -0.2157,  0.0672,  0.5696]], device='cuda:0'),
 tensor([[ 0.0603,  1.0157, -0.1446, -1.6753]], device='cuda:0'),
 tensor([[-0.0318,  0.0135,  0.0334,  0.0655]], device='cuda:0'),
 tensor([[-0.0342, -0.4375,  0.1618,  1.1126]], device='cuda:0'),
 None,
 tensor([[ 0.1095,  1.4092, -0.1383, -2.1604]], device='cuda:0'),
 tensor([[-0.0315, -0.1821,  0.0348,  0.3686]], device='cuda:0'),
 tensor([[-0.0208, -0.2307,  0.0986,  0.5520]], device='cuda:0'),
 tensor([[-0.0017,  0.5718,  0.0247, -0.6719]], device='cuda:0'),
 tensor([[-0.0562,  0.3336,  0.1494, -0.1687]], device='cuda:0'),
 None,
 tensor([[ 0.0361,  0.1762,  0.0631, -0.0539]], device='cuda:0'),
 tensor([[ 0.1826,  1.1228,  0.0472, -0.8785]], device='cuda:0'),
 tensor([[-0.0495,  0.5263,  0.1460, -0.4107]], device='cuda:0'),
 tensor([[-0.0108,  0.0153, -0.0066,  0.0259]], device='cuda:0'),
 tensor([[ 0.1062,  0.6378, -0.1170, -0.9742]], device='cuda:0'),
 tensor([[ 0.1377,  1.6054, -0.1815, -2.4924]], device='cuda:0'),
 tensor([[-0.0216, -0.1797,  0.0124,  0.3146]], device='cuda:0'),
 tensor([[-0.0243,  0.0150,  0.0069,  0.0318]], device='cuda:0'),
 tensor([[ 0.0092, -0.4291,  0.0291,  0.6114]], device='cuda:0'),
 tensor([[ 0.0806,  0.8225, -0.1781, -1.4309]], device='cuda:0'),
 tensor([[-0.0387,  0.1740, -0.0020, -0.3269]], device='cuda:0'),
 tensor([[-0.0795, -0.2135,  0.1191,  0.4876]], device='cuda:0'),
 tensor([[-4.4667e-02, -2.0743e-02,  2.7441e-05, -4.3369e-02]], device='cuda:0'),
 tensor([[ 0.0544,  0.4382, -0.0440, -0.5791]], device='cuda:0'),
 tensor([[ 0.0829,  1.0977,  0.0431, -0.9767]], device='cuda:0'),
 tensor([[ 0.1300,  1.0108, -0.2012, -1.6604]], device='cuda:0'),
 tensor([[-0.0119, -0.4301,  0.0596,  0.6337]], device='cuda:0'),
 tensor([[ 0.1379,  0.5778, -0.1564, -1.0362]], device='cuda:0'),
 tensor([[-0.0108,  0.7907,  0.0156, -1.0333]], device='cuda:0'),
 tensor([[ 0.0209,  0.0403, -0.0108, -0.0289]], device='cuda:0'),
 tensor([[-0.0580, -0.0520,  0.1416,  0.3188]], device='cuda:0'),
 tensor([[-0.0033,  0.0339, -0.0412, -0.0622]], device='cuda:0'),
 tensor([[ 0.0349, -0.0216,  0.0786,  0.2988]], device='cuda:0'),
 tensor([[ 0.0716,  0.5464,  0.1139, -0.1981]], device='cuda:0'),
 tensor([[ 0.2168,  1.1177, -0.0258, -0.8012]], device='cuda:0'),
 tensor([[-0.0316,  0.1528,  0.1350,  0.1131]], device='cuda:0'),
 tensor([[ 0.1949,  1.5791, -0.2027, -2.4921]], device='cuda:0'),
 tensor([[ 0.0407, -0.2124,  0.0488,  0.4966]], device='cuda:0'),
 tensor([[-0.0053, -0.2243,  0.0606,  0.4082]], device='cuda:0'),
 tensor([[-0.0072, -0.1800, -0.0131,  0.3227]], device='cuda:0'),
 tensor([[ 0.0197, -0.3928, -0.0185,  0.5500]], device='cuda:0'),
 tensor([[ 0.4324,  2.2959, -0.1982, -2.7246]], device='cuda:0'),
 tensor([[-0.0671, -0.2107,  0.0957,  0.4245]], device='cuda:0'),
 tensor([[ 0.2392,  1.3132, -0.0418, -1.1019]], device='cuda:0'),
 tensor([[ 0.0524,  0.5993,  0.0138, -0.8970]], device='cuda:0'),
 tensor([[ 0.1049,  0.6233, -0.0671, -0.7702]], device='cuda:0'),
 tensor([[-0.0273, -0.3930,  0.0485,  0.5547]], device='cuda:0'),
 tensor([[ 0.0232,  0.0398, -0.0158, -0.0188]], device='cuda:0'),
 tensor([[ 0.0611,  0.9604, -0.1764, -1.6490]], device='cuda:0'),
 tensor([[-0.0399, -0.0209, -0.0042, -0.0405]], device='cuda:0'),
 tensor([[ 0.2074,  1.0223, -0.2091, -1.5771]], device='cuda:0'),
 tensor([[ 0.0232,  0.3763,  0.0245, -0.5937]], device='cuda:0'),
 tensor([[ 0.1303,  0.3810, -0.1424, -0.7023]], device='cuda:0'),
 tensor([[ 0.0758,  1.5720, -0.1021, -2.2294]], device='cuda:0'),
 tensor([[-0.0078,  0.2285, -0.0344, -0.3439]], device='cuda:0'),
 tensor([[ 0.0289,  0.6130, -0.0475, -0.8798]], device='cuda:0'),
 tensor([[ 0.0127,  0.8078, -0.0242, -1.1648]], device='cuda:0'),
 tensor([[-0.0131,  0.2277, -0.0270, -0.3248]], device='cuda:0'),
 tensor([[ 0.0263,  0.2333,  0.0297, -0.1874]], device='cuda:0'),
 tensor([[ 0.0900,  0.7685, -0.0864, -1.0032]], device='cuda:0'),
 tensor([[ 0.1479,  0.9677, -0.1658, -1.3973]], device='cuda:0'),
 tensor([[ 0.1426,  1.9639, -0.1978, -2.8855]], device='cuda:0'),
 tensor([[ 0.0401,  0.2099,  0.0327, -0.3305]], device='cuda:0'))

non_final_mask

tensor([ True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True, False,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True, False,  True,
         True,  True,  True,  True, False,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True],
       device='cuda:0')

non_final_next_states

tensor([[ 1.0728e-01,  1.7679e+00, -1.4673e-01, -2.5517e+00],
        [ 8.5259e-02,  1.2131e+00, -1.0157e-01, -1.8380e+00],
        [-7.8343e-02,  5.5417e-01,  1.4926e-01, -4.0886e-01],
        [ 3.5235e-01,  1.9037e+00, -1.0908e-01, -2.0654e+00],
        [ 1.9913e-02,  1.9158e-01, -1.7274e-02, -3.0751e-01],
        [-1.9780e-02,  4.0505e-01,  2.3448e-03, -5.4902e-01],
        [ 9.7404e-03,  3.7633e-01,  1.1278e-02, -3.7153e-01],
        [-2.6455e-02,  5.4661e-01,  1.1037e-01, -2.3874e-01],
        [ 9.3618e-02,  1.0071e+00, -1.4358e-01, -1.5632e+00],
        [ 2.2926e-02,  4.2684e-01, -8.2918e-02, -7.0916e-01],
        [-3.5185e-02, -2.1130e-02, -8.5862e-03, -3.4842e-02],
        [ 3.4693e-02,  6.3463e-01, -1.4485e-01, -1.1158e+00],
        [ 1.5410e-01,  8.2174e-01, -1.3214e-01, -1.1422e+00],
        [ 1.6725e-01,  1.3828e+00, -1.5965e-01, -2.1547e+00],
        [-4.3299e-02,  4.3027e-01, -2.2846e-02, -6.0928e-01],
        [-2.8509e-02,  5.7321e-01,  5.8367e-02, -7.0439e-01],
        [-8.5573e-02,  3.6152e-01,  1.5262e-01, -1.6795e-01],
        [ 2.3138e-01,  1.5120e+00,  6.4903e-03, -1.4392e+00],
        [ 2.3984e-02, -1.5510e-01, -1.6168e-02,  2.6882e-01],
        [ 8.5098e-03,  4.0588e-02,  9.9209e-03, -3.6359e-02],
        [-6.0870e-02,  1.4486e-01,  1.4186e-01, -1.5038e-02],
        [-2.2909e-02,  7.6013e-01, -3.8622e-02, -1.2230e+00],
        [ 4.4020e-02,  9.6161e-01, -2.8407e-02, -1.2477e+00],
        [-2.5384e-02, -3.7131e-02,  1.0959e-01,  2.9197e-01],
        [ 5.7491e-02,  5.7101e-01, -2.8258e-02, -8.7609e-01],
        [-2.2680e-02, -2.0627e-01,  2.3863e-02,  3.2464e-01],
        [-1.5548e-02, -5.8772e-01,  3.1742e-02,  8.3726e-01],
        [ 2.0503e-01,  1.3173e+00,  2.9610e-02, -1.1560e+00],
        [-1.4073e-02, -3.7479e-01,  2.4043e-04,  6.0724e-01],
        [ 1.0511e-02,  1.1017e+00,  1.0213e-01, -1.0711e+00],
        [-2.6127e-02, -2.3363e-01,  1.1543e-01,  6.1711e-01],
        [ 6.1860e-04, -6.2460e-01,  4.1290e-02,  9.1309e-01],
        [ 3.7959e-02, -2.3878e-02,  8.5210e-02,  3.5006e-01],
        [ 2.6300e-02, -1.5563e-01, -2.1402e-02,  2.8052e-01],
        [-4.4911e-03,  2.7487e-02, -6.8719e-03,  1.8818e-03],
        [ 3.5180e-02, -2.0672e-01,  3.9325e-02,  3.7091e-01],
        [ 5.1980e-02,  8.1759e-01,  1.4439e-03, -1.0424e+00],
        [ 5.0652e-02,  7.0924e-01,  6.5745e-02, -4.3127e-01],
        [ 4.9577e-02,  2.4258e-01, -3.8557e-02, -2.7447e-01],
        [ 3.6487e-02, -1.8007e-02,  5.8725e-02,  2.1971e-01],
        [-4.1594e-02,  3.6951e-01, -7.5608e-03, -6.2899e-01],
        [-2.0081e-02, -3.4411e-02,  9.3927e-02,  2.3128e-01],
        [ 5.0384e-02,  5.3801e-01,  7.4328e-02, -4.7782e-02],
        [-8.8230e-03,  4.3273e-01, -7.2807e-02, -6.6481e-01],
        [-3.4912e-02,  2.0700e-01,  4.3867e-02, -1.9207e-01],
        [-8.8472e-02, -2.3938e-02,  1.4481e-01,  3.1723e-01],
        [-2.2179e-02,  4.3108e-01, -5.3213e-02, -6.2760e-01],
        [ 3.2192e-01,  1.5108e+00, -1.1469e-01, -1.4554e+00],
        [ 6.1144e-02,  7.3199e-01,  7.3373e-02, -3.1612e-01],
        [ 1.9058e-01,  1.3128e+00, -3.9148e-03, -1.0927e+00],
        [-3.2662e-02, -1.8011e-01,  3.0902e-02,  3.2445e-01],
        [-1.0478e-02, -1.7975e-01, -6.0892e-03,  3.1648e-01],
        [ 4.7386e-02,  8.3133e-01, -1.6717e-01, -1.4501e+00],
        [ 2.3502e-03,  7.3699e-01,  1.0391e-01, -4.2796e-01],
        [-2.4166e-04, -2.2275e-01,  5.1142e-02,  3.7374e-01],
        [-4.6394e-02, -4.3457e-01,  1.1288e-01,  7.3644e-01],
        [ 6.8068e-02,  4.3962e-01, -6.1639e-02, -6.1027e-01],
        [-7.2920e-03,  2.7683e-02, -9.8648e-04, -2.4419e-03],
        [ 2.6546e-01,  1.5088e+00, -6.3831e-02, -1.4074e+00],
        [-5.2322e-02,  5.5012e-01,  1.2807e-01, -3.1744e-01],
        [-2.6805e-02, -4.0172e-01,  3.0356e-02,  6.2475e-01],
        [-2.5162e-02, -3.7496e-01,  1.8678e-02,  6.1120e-01],
        [ 3.7983e-02,  1.1537e+00, -1.3846e-01, -1.8957e+00],
        [ 3.9257e-02, -2.1569e-01,  6.7202e-02,  5.6961e-01],
        [ 6.0309e-02,  1.0157e+00, -1.4460e-01, -1.6753e+00],
        [-3.1790e-02,  1.3482e-02,  3.3445e-02,  6.5509e-02],
        [-3.4244e-02, -4.3754e-01,  1.6175e-01,  1.1126e+00],
        [ 1.0952e-01,  1.4092e+00, -1.3833e-01, -2.1604e+00],
        [-3.1520e-02, -1.8210e-01,  3.4755e-02,  3.6855e-01],
        [-2.0769e-02, -2.3074e-01,  9.8553e-02,  5.5205e-01],
        [-1.6954e-03,  5.7179e-01,  2.4716e-02, -6.7189e-01],
        [-5.6196e-02,  3.3359e-01,  1.4941e-01, -1.6868e-01],
        [ 3.6127e-02,  1.7623e-01,  6.3119e-02, -5.3885e-02],
        [ 1.8258e-01,  1.1228e+00,  4.7180e-02, -8.7851e-01],
        [-4.9524e-02,  5.2630e-01,  1.4604e-01, -4.1075e-01],
        [-1.0784e-02,  1.5276e-02, -6.6071e-03,  2.5892e-02],
        [ 1.0620e-01,  6.3776e-01, -1.1701e-01, -9.7425e-01],
        [ 1.3770e-01,  1.6054e+00, -1.8154e-01, -2.4924e+00],
        [-2.1569e-02, -1.7967e-01,  1.2385e-02,  3.1463e-01],
        [-2.4281e-02,  1.5009e-02,  6.8833e-03,  3.1789e-02],
        [ 9.2003e-03, -4.2908e-01,  2.9062e-02,  6.1140e-01],
        [ 8.0622e-02,  8.2249e-01, -1.7811e-01, -1.4309e+00],
        [-3.8664e-02,  1.7396e-01, -2.0487e-03, -3.2688e-01],
        [-7.9457e-02, -2.1351e-01,  1.1906e-01,  4.8758e-01],
        [-4.4667e-02, -2.0743e-02,  2.7441e-05, -4.3369e-02],
        [ 5.4428e-02,  4.3823e-01, -4.4046e-02, -5.7906e-01],
        [ 8.2904e-02,  1.0977e+00,  4.3069e-02, -9.7669e-01],
        [ 1.3004e-01,  1.0108e+00, -2.0121e-01, -1.6604e+00],
        [-1.1873e-02, -4.3006e-01,  5.9552e-02,  6.3367e-01],
        [ 1.3795e-01,  5.7777e-01, -1.5642e-01, -1.0362e+00],
        [-1.0823e-02,  7.9069e-01,  1.5605e-02, -1.0333e+00],
        [ 2.0882e-02,  4.0251e-02, -1.0792e-02, -2.8917e-02],
        [-5.7973e-02, -5.1984e-02,  1.4156e-01,  3.1882e-01],
        [-3.2703e-03,  3.3895e-02, -4.1228e-02, -6.2195e-02],
        [ 3.4943e-02, -2.1571e-02,  7.8594e-02,  2.9884e-01],
        [ 7.1582e-02,  5.4641e-01,  1.1389e-01, -1.9807e-01],
        [ 2.1684e-01,  1.1177e+00, -2.5768e-02, -8.0122e-01],
        [-3.1605e-02,  1.5280e-01,  1.3503e-01,  1.1308e-01],
        [ 1.9491e-01,  1.5791e+00, -2.0274e-01, -2.4921e+00],
        [ 4.0735e-02, -2.1241e-01,  4.8793e-02,  4.9663e-01],
        [-5.2644e-03, -2.2430e-01,  6.0569e-02,  4.0820e-01],
        [-7.1832e-03, -1.8003e-01, -1.3060e-02,  3.2267e-01],
        [ 1.9717e-02, -3.9284e-01, -1.8533e-02,  5.4995e-01],
        [ 4.3242e-01,  2.2959e+00, -1.9818e-01, -2.7246e+00],
        [-6.7102e-02, -2.1072e-01,  9.5657e-02,  4.2452e-01],
        [ 2.3919e-01,  1.3132e+00, -4.1793e-02, -1.1019e+00],
        [ 5.2352e-02,  5.9933e-01,  1.3845e-02, -8.9700e-01],
        [ 1.0494e-01,  6.2334e-01, -6.7060e-02, -7.7024e-01],
        [-2.7302e-02, -3.9304e-01,  4.8487e-02,  5.5473e-01],
        [ 2.3188e-02,  3.9794e-02, -1.5792e-02, -1.8837e-02],
        [ 6.1058e-02,  9.6035e-01, -1.7637e-01, -1.6490e+00],
        [-3.9930e-02, -2.0874e-02, -4.1807e-03, -4.0477e-02],
        [ 2.0741e-01,  1.0223e+00, -2.0911e-01, -1.5771e+00],
        [ 2.3221e-02,  3.7635e-01,  2.4532e-02, -5.9369e-01],
        [ 1.3033e-01,  3.8099e-01, -1.4237e-01, -7.0231e-01],
        [ 7.5844e-02,  1.5720e+00, -1.0214e-01, -2.2294e+00],
        [-7.8406e-03,  2.2851e-01, -3.4351e-02, -3.4385e-01],
        [ 2.8878e-02,  6.1299e-01, -4.7538e-02, -8.7981e-01],
        [ 1.2722e-02,  8.0779e-01, -2.4242e-02, -1.1648e+00],
        [-1.3052e-02,  2.2765e-01, -2.7039e-02, -3.2481e-01],
        [ 2.6301e-02,  2.3327e-01,  2.9704e-02, -1.8744e-01],
        [ 9.0040e-02,  7.6855e-01, -8.6415e-02, -1.0032e+00],
        [ 1.4793e-01,  9.6767e-01, -1.6583e-01, -1.3973e+00],
        [ 1.4264e-01,  1.9639e+00, -1.9776e-01, -2.8855e+00],
        [ 4.0061e-02,  2.0994e-01,  3.2707e-02, -3.3046e-01]], device='cuda:0')

state_batch.shape

torch.Size([128, 4])

action_batch.shape

torch.Size([128, 1])

reward_batch

tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
        1., 1.], device='cuda:0')

__t

tensor([[0.8239, 1.0678],
        [1.5102, 1.8625],
        [1.1531, 1.4095],
        [0.7606, 1.0270],
        [0.7679, 1.0549],
        [0.7800, 0.9981],
        [0.9312, 1.3687],
        [1.0128, 1.3304],
        [1.6723, 2.1253],
        [0.7289, 0.9694],
        [1.7867, 2.2589],
        [0.7771, 1.0487],
        [1.1555, 1.4792],
        [1.4739, 1.8754],
        [0.7970, 1.0090],
        [0.8166, 1.0341],
        [0.7520, 0.9775],
        [1.4251, 1.8186],
        [0.7406, 0.9626],
        [0.8206, 1.0935],
        [1.2606, 1.6089],
        [0.9438, 1.2157],
        [0.7455, 0.9680],
        [1.5628, 1.9837],
        [0.8967, 1.1652],
        [0.7654, 0.9682],
        [0.8783, 1.1400],
        [0.7707, 1.0308],
        [1.4016, 1.7157],
        [0.7585, 1.0263],
        [0.7604, 1.0498],
        [0.7715, 1.0034],
        [0.7650, 1.0434],
        [1.0396, 1.2708],
        [1.2695, 1.6362],
        [1.3621, 1.7203],
        [1.4564, 1.8477],
        [0.7637, 1.0681],
        [0.7436, 0.9637],
        [0.8141, 1.0474],
        [0.7496, 0.9662],
        [1.0074, 1.3075],
        [0.8058, 1.1415],
        [1.0222, 1.3213],
        [0.7964, 1.0206],
        [0.7444, 0.9654],
        [1.0712, 1.4032],
        [0.8287, 1.0278],
        [0.7791, 1.0071],
        [1.3797, 1.7582],
        [0.8031, 1.1290],
        [1.1220, 1.4722],
        [0.8581, 1.1328],
        [0.7458, 0.9692],
        [1.9925, 2.5241],
        [0.8844, 1.1042],
        [0.7859, 1.0129],
        [0.7419, 0.9734],
        [1.0514, 1.3672],
        [1.2441, 1.5828],
        [0.7758, 1.0221],
        [0.7502, 0.9799],
        [0.7625, 0.9714],
        [1.3229, 1.7080],
        [0.7379, 0.9676],
        [1.1362, 1.3960],
        [0.7507, 0.9510],
        [1.3220, 1.6228],
        [0.7775, 1.0539],
        [0.8222, 1.1699],
        [1.6477, 2.1127],
        [1.0839, 1.3923],
        [0.8092, 1.1518],
        [0.8763, 1.0864],
        [1.1191, 1.4253],
        [1.0011, 1.2350],
        [1.3038, 1.6590],
        [1.3925, 1.7687],
        [1.1646, 1.4218],
        [0.7395, 0.9621],
        [0.7741, 1.0439],
        [0.9269, 1.1422],
        [1.1394, 1.3959],
        [1.2993, 1.6673],
        [0.7799, 1.0119],
        [0.7707, 1.0735],
        [2.0438, 2.5919],
        [0.9913, 1.2691],
        [0.7499, 0.9620],
        [0.9051, 1.3184],
        [0.9155, 1.1439],
        [0.7461, 0.9607],
        [0.7499, 0.9686],
        [0.7532, 0.9851],
        [1.9530, 2.4456],
        [1.0390, 1.5443],
        [0.7440, 0.9491],
        [1.2709, 1.6259],
        [0.7590, 1.0258],
        [0.7634, 1.0303],
        [0.7521, 0.9729],
        [1.0305, 1.2570],
        [1.7204, 2.1211],
        [0.7467, 0.9588],
        [0.8260, 1.0523],
        [0.9137, 1.2058],
        [0.7448, 0.9628],
        [1.1614, 1.4941],
        [0.8785, 1.0930],
        [0.9300, 1.2236],
        [1.3657, 1.7324],
        [0.9146, 1.1254],
        [1.2236, 1.5748],
        [0.7643, 1.0162],
        [0.7571, 1.0082],
        [0.8118, 1.0476],
        [1.8991, 2.3756],
        [0.7402, 0.9630],
        [0.7951, 1.0291],
        [0.8198, 1.1739],
        [1.1728, 1.4595],
        [1.5033, 1.9224],
        [0.9975, 1.3053],
        [1.1705, 1.5177],
        [0.7491, 0.9692],
        [0.7720, 1.0358],
        [0.7649, 0.9644],
        [0.7320, 0.9879]], device='cuda:0', grad_fn=<AddmmBackward0>)

expected_state_action_values

tensor([2.2928, 1.9625, 1.4207, 2.1499, 1.4084, 1.4506, 1.4171, 1.4036, 1.8529,
        1.5080, 1.3836, 1.6617, 1.6917, 2.1036, 1.4724, 1.4929, 1.3812, 1.8639,
        1.4004, 1.3794, 1.3526, 1.6855, 1.7178, 1.3750, 1.5543, 1.0000, 1.4032,
        1.5156, 1.7385, 1.4519, 1.6601, 1.4305, 1.5354, 1.3849, 1.4020, 1.3803,
        1.4091, 1.6270, 1.4476, 1.4067, 1.3784, 1.4689, 1.3721, 1.3883, 1.4954,
        1.3855, 1.3683, 1.4826, 1.8911, 1.4277, 1.7251, 1.3977, 1.4022, 1.7976,
        1.4444, 1.4073, 1.4701, 1.4811, 1.3793, 1.8667, 1.4113, 1.4553, 1.4508,
        1.9791, 1.4284, 1.8901, 1.3705, 1.5423, 1.0000, 2.1070, 1.4011, 1.4223,
        1.4908, 1.3795, 1.0000, 1.3712, 1.6187, 1.4180, 1.3797, 1.6121, 2.2511,
        1.3996, 1.3766, 1.4590, 1.7936, 1.4091, 1.4084, 1.3822, 1.4698, 1.6433,
        1.9002, 1.4589, 1.6349, 1.6177, 1.3832, 1.3732, 1.3877, 1.3816, 1.3966,
        1.6092, 1.3519, 2.2492, 1.4199, 1.4088, 1.4035, 1.4505, 2.4448, 1.4024,
        1.7346, 1.5550, 1.5385, 1.4438, 1.3838, 1.8823, 1.3829, 1.8780, 1.4538,
        1.5172, 2.1501, 1.4165, 1.5649, 1.6701, 1.4123, 1.3880, 1.6303, 1.7996,
        2.4402, 1.4056], device='cuda:0')